#include <gtest/gtest.h>
#include "Common.h"
#include <omp.h>
#include <numeric>
#include "WorkerTool.h"

#if defined(USE_NPU)
    #include "acl/acl.h"
#elif defined(USE_GPU)
    #include <cuda_runtime.h>
#endif

class KVCacheWorkerToolTest : public ::testing::Test {
protected:
  void SetUp() override {
  }

  void TearDown() override {
  }
};

#define BLOCK_SIZE (16 * 1024 * 1024)

void fillSegmentArgs(int32_t blkNum, std::map<int32_t, Segment>& segMap)
{
    for (int32_t i = 0; i < blkNum; i++) {
        void *memAddr = malloc(BLOCK_SIZE);
        assert(memAddr != nullptr);
        segMap[i] = Segment(i, i * BLOCK_SIZE, (uint8_t *)memAddr, BLOCK_SIZE, false, "", 1, 0, 63, 64, BLOCK_SIZE / 64, i, i);
    }
}

void freeSegmentArgs(int32_t blkNum, std::map<int32_t, Segment>& segMap)
{
    for (int32_t i = 0; i < blkNum; i++) {
        uint8_t *memAddr = segMap[i].getMemAddr();
        free(memAddr);
    }
}

TEST_F(KVCacheWorkerToolTest, TestLayerBuf) {
    auto tool = WorkerTool(DEVICE_TYPE_CPU);
    int32_t blkNum = 10;

    std::map<int32_t, Segment> readSegMap;
    std::map<int32_t, Segment> writeSegMap;
    fillSegmentArgs(blkNum, writeSegMap);
    fillSegmentArgs(blkNum, readSegMap);

    int32_t ret = tool.updateLayerBufInfo(readSegMap, writeSegMap, 0, 128 * 1024, 8, 128, 2, true, 0, 1);
    assert(ret == KVC_OK);

    std::map<int32_t, Segment> decodeSegMap;
    ret = tool.adjustUpdateLayerBufInfo(readSegMap, writeSegMap, decodeSegMap, 0, 128 * 1024, 8, 128, 2, true, 0, 1);
    assert(ret == KVC_OK);

    freeSegmentArgs(blkNum, writeSegMap);
    freeSegmentArgs(blkNum, readSegMap);
}

TEST_F(KVCacheWorkerToolTest, TestSwapByLayerContinuous) {
    int32_t blockBytes = 128 * 1024;
    CacheInfo cacheInfo{ 0, 63, 64, 256 * 1024};

    int32_t blkNum = 10;
    std::vector<int32_t> phyBlkIds;
    std::vector<Segment> segments;
    for (int32_t i = 0; i < blkNum; i++) {
        phyBlkIds.emplace_back(i);
        void *memAddr = malloc(BLOCK_SIZE);
        assert(memAddr != nullptr);
        segments.emplace_back(i, i * BLOCK_SIZE, (uint8_t *)memAddr, BLOCK_SIZE, false, "", 1, 0, 63, 64, BLOCK_SIZE / 64, i, i);
    }

    // 2 * num_blocks * block_size * num_heads * head_size * element_size;
    int64_t layer_cache_size = 2LL * blkNum * 64 * 8 * 128 * 2;
    char *testBuf = (char*)malloc(layer_cache_size);
    assert(testBuf != nullptr);
    std::fill(testBuf, testBuf + layer_cache_size, 49);
    void *layer_cache_ptr;
    DeviceType type;
#if defined(USE_GPU)
    type = DEVICE_TYPE_GPU;
    cudaMalloc(&layer_cache_ptr, layer_cache_size);
    assert(layer_cache_ptr != nullptr);
    cudaMemcpy(layer_cache_ptr, (void *)testBuf, layer_cache_size, cudaMemcpyHostToDevice);
#elif defined(USE_NPU)
    type = DEVICE_TYPE_NPU;
    assert(0 == aclrtSetDevice(0));
    assert(0 == aclrtMalloc(&layer_cache_ptr, layer_cache_size, ACL_MEM_MALLOC_NORMAL_ONLY));
    assert(nullptr != layer_cache_ptr);
    aclrtMemcpy(layer_cache_ptr, layer_cache_size, (char *)testBuf, layer_cache_size, ACL_MEMCPY_HOST_TO_DEVICE);
#else
    type = DEVICE_TYPE_CPU;
    layer_cache_ptr = (void *)testBuf;
#endif

    LayCacheInfo info = {
        .layerIdx = 0,
        .numBlocks = blkNum,
        .transferUnitSize = blockBytes,
        .hasValue = true,
        .tpNum = 1,
        .tpRank = 0,
        .stream = nullptr,
    };
    std::cout << "phyBlkIds num: " << phyBlkIds.size() << std::endl;
    uint64_t start = getCurrentTimestampUs();
    swapOneLayerWithoutTP(segments, layer_cache_ptr, phyBlkIds, DIRECTION::DEVICE_TO_HOST, type, info);
    // std::cout << "[write to dram without tp] time cost: " << getCurrentTimestampUs() - start << " us" << std::endl;
    for (int32_t i = 0; i < blkNum; i++) {
        uint8_t *memAddr = segments[i].getMemAddr();
        for (int32_t j = 0; j < blockBytes; j++) {
            assert(memAddr[j] == 49);
        }
        free(memAddr);
    }

#if defined(USE_GPU)
    cudaFree(layer_cache_ptr);
#elif defined(USE_NPU)
    aclrtFree(layer_cache_ptr);
#endif
    free(testBuf);
    return;
}

int main(int argc, char** argv) {
    testing::InitGoogleTest(&argc, argv);
    return RUN_ALL_TESTS();
}
